from mitmproxy import proxy, options
from mitmproxy.tools.dump import DumpMaster
from mitmproxy import http, websocket
import logging


class WebSocketHandler:
    def websocket_start(self, flow: websocket.WebSocketFlow):
        '''
        WebSocket 连接开始时被调用。
        :param flow: 当前的 WebSocket 流
        :return:
        '''
        print("WebSocket 连接开启！")
        print(f"  客户端地址: {flow.client_conn.address}")
        print(f"  服务器主机名: {flow.request.host}")
        print(f"  WebSocket请求路径: {flow.request.path}")

    def websocket_message(self, flow: websocket.WebSocketFlow):
        '''
        WebSocket 消息接收时被调用。
        :param flow: 当前的 WebSocket 流
        :return:
        '''
        # 判断最后一条消息是否来自客户端
        if flow.messages[-1].from_client:
            # 如果消息来自客户端：
            print("Client -> Server:")
        else:
            # 如果消息来自服务端
            print("Server -> Client:")
        # 输出消息内容
        print(f"  Message: {flow.messages[-1].content}")

    def websocket_end(self, flow: websocket.WebSocketFlow):
        '''
        表示WebSocket连接关闭的代码
        :param flow:
        :return:
        '''
        print("WebSocket 连接关闭！")
        print(f"  Close code: {flow.close_code}")
        print(f"  Close reason: {flow.close_message}")


addons = [
    WebSocketHandler(),
]


def start_mitmproxy():
    opts = options.Options(listen_host="127.0.0.1", listen_port=6666)
    # 设置日志级别为 info
    logging.basicConfig(level=logging.INFO)

    pconf = proxy.config.ProxyConfig(opts)
    m = DumpMaster(opts)
    m.addons.add(*addons)
    m.server = proxy.server.ProxyServer(pconf)
    m.run()


if __name__ == "__main__":
    # logging.info("开始监听 WebSocket 消息...")
    print("开始监听 WebSocket 消息...")
    start_mitmproxy()
